package middleware

import (
	"context"
	"log"
	"oc_saas/glo"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/go-redis/redis/v8"
)

// 限流器
// 令牌桶限流器结构
type RateLimiter struct {
	RedisClient *redis.Client
	Key         string        //存储令牌桶信息的键的名称
	MaxTokens   int           //最大令牌数
	FillRate    time.Duration //令牌桶填充速率的时间间隔
}

// NewRateLimiter 创建一个 RateLimiter 实例
func NewRateLimiter(redisClient *redis.Client, key string, maxTokens int, fillRate time.Duration) *RateLimiter {
	return &RateLimiter{
		RedisClient: redisClient,
		Key:         key,
		MaxTokens:   maxTokens,
		FillRate:    fillRate,
	}
}

// Allow 检查是否允许请求通过
func (l *RateLimiter) Allow() bool {
	// 获取当前令牌数量
	count, err := l.RedisClient.Get(context.Background(), l.Key).Int()
	if err != nil {
		log.Fatal(err)
		return false
	}
	if count <= 0 {
		return false
	}
	l.RedisClient.Decr(context.Background(), l.Key)
	return true
}

// FillBucket 填充令牌桶
func (l *RateLimiter) FillBucket() {
	// 获取当前令牌数量
	count, err := l.RedisClient.Get(context.Background(), l.Key).Int()
	if err != nil && err != redis.Nil {
		log.Fatal(err)
		return
	}
	// 计算要增加的令牌数量
	addTokens := l.MaxTokens - count
	if addTokens > 0 {
		l.RedisClient.IncrBy(context.Background(), l.Key, int64(addTokens))
	}
}

// 中间件
func RateLimiterMiddleware(limit int, key string, fileRate time.Duration) gin.HandlerFunc {
	l := NewRateLimiter(glo.OC_REDIS, key, limit, fileRate)
	go func() {
		ticker := time.NewTicker(fileRate)
		for range ticker.C {
			l.FillBucket()
		}
	}()
	return func(c *gin.Context) {
		if !l.Allow() {
			return
		}
		c.Next()
	}
}
